#include "stdafx.h"
#include "SimpleSocket.h"
#include "WindowsUtility.h"

namespace RoomFW {

// Make sure if you enable G_TRACE_METHOD, that you've initialized the logger in your setup!
#if 0		// DEBUG -- hansohn, 15 May 2009
	#define TRACE_ENTRY G_TRACE_METHOD
#else
	#define TRACE_ENTRY(x)
#endif

SimpleSocket::SimpleSocket(void) :
	socket_(INVALID_SOCKET),
	isListener_(false)
{
}

SimpleSocket::SimpleSocket(SOCKET socket) :
	socket_(socket),
	isListener_(false)
{
G_LOG_1_TRACE("SimpleSocket::SimpleSocket socket_:%d", socket_);
}

SimpleSocket::~SimpleSocket()
{
	this->close(!isListener_);
}

SimpleSocket::operator SOCKET() const
{
	return socket_;
}

bool SimpleSocket::valid(void) const
{
	return (INVALID_SOCKET != socket_);
}

bool SimpleSocket::listen(unsigned short port, const char* ipAddress)
{
	static const char* functionName = "SimpleSocket::listen";
	
	TRACE_ENTRY(SimpleSocket::listen)
	
	ASSERT(!this->valid());

	bool retVal = false;

	// Create the socket.
	socket_ = ::socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	SHOW_SOCKET_ERROR((INVALID_SOCKET != socket_), functionName, "::socket");
	if (INVALID_SOCKET != socket_)
	{
		// Bind socket to address.
		struct sockaddr_in socketAddress = InitializeSockaddr(port, ipAddress);
		int ret = ::bind(socket_, (SOCKADDR*)&socketAddress, sizeof(socketAddress));
		SHOW_SOCKET_ERROR((SOCKET_ERROR != ret), functionName, "::bind");
		if (SOCKET_ERROR != ret)
		{
			// Start listening (default backlog).
			ret = ::listen(socket_, SOMAXCONN);
			SHOW_SOCKET_ERROR((SOCKET_ERROR != ret), functionName, "::listen");
			if (SOCKET_ERROR != ret)
			{
				// Set to non-blocking mode.
				u_long ignored;
				if (this->ioclt(FIONBIO, &ignored))
				{
					retVal = true;
					isListener_ = true;
				}
			}
		}
		if (!retVal)
		{
			this->close(false);
		}
	}
	if (retVal)
		G_LOG_4_TRACE("%s(%d, %s), socket_:%d", functionName, port, ipAddress, socket_);
	
	return retVal;
}

SimpleSocket::SocketPtr SimpleSocket::accept(void)
{
	static const char* functionName = "SimpleSocket::accept";
	
	TRACE_ENTRY(SimpleSocket::SocketPtr SimpleSocket::accept);
	
	ASSERT(this->valid());

	// Accept the connection.  Don't needs its address.
	SocketPtr retVal;
	SOCKET retSock = ::accept(socket_, 0, 0);
	SHOW_SOCKET_ERROR((INVALID_SOCKET != retSock), functionName, "::accept");
	if (INVALID_SOCKET != retSock)
	{
		// Return a new simple socket smart pointer.
		retVal = SocketPtr(new SimpleSocket(retSock));
	}
	return retVal;
}

bool SimpleSocket::connect(unsigned short port, const char* ipAddress)
{
	static const char* functionName = "SimpleSocket::connect";
	
	ASSERT(!this->valid());

	bool retVal = false;

	// Create the socket.
	socket_ = ::socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	SHOW_SOCKET_ERROR((INVALID_SOCKET != socket_), functionName, "::socket");
	if (INVALID_SOCKET != socket_)
	{
		// Connect to address.
		struct sockaddr_in socketAddress = InitializeSockaddr(port, ipAddress);
		int ret = ::connect(socket_, (SOCKADDR*)&socketAddress, sizeof(socketAddress));
		SHOW_SOCKET_ERROR((SOCKET_ERROR != ret), functionName, "::connect");
		if (SOCKET_ERROR != ret)
		{
			// Set to non-blocking mode.
			u_long ignored;
			if (this->ioclt(FIONBIO, &ignored))
			{
				retVal = true;
				isListener_ = true;
			}
		}
		else
		{
			this->close(false);
		}
	}
	G_LOG_5_TRACE("%s(%d, %s), socket_:%d %s", functionName, port, ipAddress, socket_, (retVal ? "" : "FAILED"));
	return retVal;
}

bool SimpleSocket::send(const char* buffer, int size) const
{
	static const char* functionName = "SimpleSocket::send";
	
	TRACE_ENTRY(SimpleSocket::send);
	
	ASSERT(this->valid());
	ASSERT((0 != buffer) || (0 == size));
	ASSERT(size >= 0);

	bool retVal = this->waitWrite();
	int bytesSent = 0;
	if (retVal) do
	{
		G_LOG_3_TRACE("%s ::send(%d, ..., %d)", functionName, socket_, size);
		bytesSent = ::send(socket_, buffer, size, 0);
		SHOW_SOCKET_ERROR((SOCKET_ERROR != bytesSent), functionName, "::send");
		if (SOCKET_ERROR != bytesSent)
		{
			buffer += bytesSent;
			size -= bytesSent;
		}
		else
		{
			retVal = false;
		}
	} while ((bytesSent > 0) && (size > 0) && this->waitWrite());
	return retVal;
}

bool SimpleSocket::sendMessage(const std::string& message) const
{
	bool retVal = false;

	if(this->valid())
	{
		// First send message size.
		unsigned long messageSize = htonl(static_cast<unsigned long>(message.size()));
		int byteCount = sizeof(messageSize);
		if (this->send(reinterpret_cast<const char*>(&messageSize), byteCount))
		{
			// Send message.
			byteCount = message.size();
			if (this->send(message.data(), byteCount))
			{
				retVal = true;
			}
		}
	}
	else
	{
		G_LOG_0_WARNING("SimpleSocket::sendMessage(): invalid socket");
	}
	return retVal;
}

int SimpleSocket::recv(SOCKET s, char* buf, int len, int flags) const
{
#ifdef _DONT_LOG_PEEKS_
	if (flags == MSG_PEEK)
	{
		return ::recv(s, buf, len, flags);
	}
#endif

	int retVal = ::recv(s, buf, len, flags);
	return retVal;
}

int SimpleSocket::select(int nfds, fd_set* readfds, fd_set* writefds, fd_set* exceptfds, const struct timeval* timeout) const
{
	static const char* functionName = "SimpleSocket::select";
	
	TRACE_ENTRY(SimpleSocket::select);

	G_LOG_5_TRACE("%s(%d, %s, %s) %s", functionName,
		nfds, (readfds ? "readfds" : ""), (writefds ? "writefds" : ""), (timeout ? "" : "BLOCKING"));

	int result = ::select(nfds, readfds, writefds, exceptfds, timeout);
	return result;
}

bool SimpleSocket::receive(char* buffer, int& size) const
{
	static const char* functionName = "SimpleSocket::receive";
	
	TRACE_ENTRY(SimpleSocket::receive);
	
	ASSERT(this->valid());
	ASSERT(0 != buffer);
	ASSERT(size > 0);

	bool retVal = this->waitRead();
	int bytesReceived = 0;
	if (!retVal) 
	{
		G_LOG_1_TRACE("%s:  waitRead() returned false", functionName);
	}
	else
	{
		bool done = false;
		do
		{
			bytesReceived = recv(socket_, buffer, size, 0);
			if (SOCKET_ERROR == bytesReceived)
			{
				// If connection reset by remote host, just consider it closed, not error.
				int lastSocketError = WSAGetLastError();
				switch(lastSocketError)
				{
				case WSAEWOULDBLOCK:
					SHOW_SOCKET_ERROR(false, functionName, "retrying recv");
					break;
				case WSAECONNRESET:
					done = true;
					break;		// Ignore these -- they can happen if the socket has been closed.
				default:
					SHOW_SOCKET_ERROR(false, functionName, "recv");
					done = true;
					retVal = false;
				}
			}
			else if (0 == bytesReceived )
			{
				// http://msdn.microsoft.com/en-us/library/ms740121(VS.85).aspx
				G_LOG_1_TRACE("%s: call to recv returned 0 indicating socket gracefully closed", functionName);
				done = true;
			}			
			else
			{
				buffer += bytesReceived;
				size -= bytesReceived;
				ASSERT(size >= 0);
				if (size <= 0)
					done = true;
			}
		} while (!done && this->waitRead());
	}

	if (size > 0)
		G_LOG_2_TRACE("%s, didn't get %d bytes expected", functionName, size);

	return retVal;
}

bool SimpleSocket::receiveMessage(std::string& message) const
{
	static const char* functionName = "SimpleSocket::receiveMessage";
	
	ASSERT(this->valid());

	bool retVal = false;

	// First get message size.
	unsigned long messageSize = 0;
	int byteCount = sizeof(messageSize);

	CriticalSectionGuard g(cs_);		// Message size and payload must be atomically received.

	if (this->receive(reinterpret_cast<char*>(&messageSize), byteCount))
	{
		// True return with bytes left indicates socket closed.
		if (byteCount > 0)
		{
			message.clear();
			retVal = true;
		}
		else // Get message
		{
			messageSize = ntohl(messageSize);
			message.resize(messageSize);

			G_LOG_3_TRACE("%s, received messageSize: %d on socket_:%d", functionName, messageSize, socket_);

			byteCount = messageSize;
			int payloadSize = byteCount;
			if (this->receive(const_cast<char*>(message.data()), byteCount))
			{
				// True return with bytes left indicates socket closed.
				if (byteCount > 0)
				{
					message.clear();
				}
				else
				{
					G_LOG_3_TRACE("%s, received %d payload bytes on socket_:%d", functionName, payloadSize, socket_);
				}
				retVal = true;
			}
		}
	}
	return retVal;
}

bool SimpleSocket::close(bool doShutdown)
{
	static const char* functionName = "SimpleSocket::close";
	
	bool retVal = !this->valid();
	if (!retVal)
	{
		int ret = SOCKET_ERROR;
		if (doShutdown)
		{
			ret = ::shutdown(socket_, SD_BOTH);
			SHOW_SOCKET_ERROR((SOCKET_ERROR != ret), functionName, "::shutdown");
		}
		ret = ::closesocket(socket_); // Base return value on close success.
		/// Note:  (Sometimes?) getting "Error 10093: WSAStartup not performed" -- See http://www.talkbass.com/forum/archive/index.php/t-161505.html
		SHOW_SOCKET_ERROR((SOCKET_ERROR != ret), functionName, "::closesocket");
		if (SOCKET_ERROR != ret)
		{
			retVal = true;
		}
		G_LOG_3_TRACE("%s, socket_:%d %s", functionName, socket_, (retVal ? "" : "FAILED"));
		socket_ = INVALID_SOCKET;
		isListener_ = false;
	}
	return retVal;
}

bool SimpleSocket::isWritable(void) const
{
	timeval timeout = {0, 0};
	return waitWrite(&timeout);
}

bool SimpleSocket::waitWrite(const timeval* timeout) const
{
	static const char* functionName = "SimpleSocket::waitWrite";
	
	ASSERT(this->valid());

	bool retVal = false;

	// We need to ensure that the socket hasn't already been closed by the other end,
	// otherwise, select won't return.
	char c;
	int bytesReady = recv(socket_, &c, 1, MSG_PEEK);
	if ((1 == bytesReady) || ((SOCKET_ERROR == bytesReady) && (WSAEWOULDBLOCK == WSAGetLastError())))
	{
		// Set up and call select.
		fd_set fdSet;
		FD_ZERO(&fdSet);
#pragma warning( suppress : 4127 )
		FD_SET(socket_, &fdSet);
		int ret = select(0, 0, &fdSet, 0, timeout);
		SHOW_SOCKET_ERROR((SOCKET_ERROR != ret), functionName, "::select");
		if (SOCKET_ERROR != ret)
		{
			retVal = ((1 == fdSet.fd_count) && (fdSet.fd_array[0] == socket_));
		}
	}
	return retVal;
}

bool SimpleSocket::isReadable(void) const
{
	timeval timeout = {0, 0};
	return waitRead(&timeout);
}

bool SimpleSocket::waitRead(const timeval* timeout) const
{
	static const char* functionName = "SimpleSocket::waitRead";
	
	TRACE_ENTRY(SimpleSocket::waitRead);
	
	ASSERT(this->valid());

	bool retVal = false;

	// We need to ensure that the socket hasn't already been closed by the other end,
	// otherwise, select won't return.
	char c;
	int len = sizeof(c);
	int bytesReady = recv(socket_, &c, len, MSG_PEEK);
	if (len == bytesReady)
	{
		retVal = true;
	}
	else if (0 == bytesReady)
	{
		// See http://msdn.microsoft.com/en-us/library/ms740121(VS.85).aspx
		G_LOG_1_TRACE("%s: call to recv returned 0 indicating socket gracefully closed", functionName);
	}
	else if ((SOCKET_ERROR == bytesReady) && (WSAEWOULDBLOCK == ::WSAGetLastError()))
	{
		// Set up and call select.
		fd_set fdSet;
		FD_ZERO(&fdSet);
#pragma warning( suppress : 4127 )
		FD_SET(socket_, &fdSet);
		int ret = select(0, &fdSet, 0, 0, timeout);
		SHOW_SOCKET_ERROR((SOCKET_ERROR != ret), functionName, "::select");
		if (SOCKET_ERROR != ret)
		{
			retVal = ((1 == fdSet.fd_count) && (fdSet.fd_array[0] == socket_));
		}
	}
	G_LOG_2_TRACE("%s: returned %d", functionName, retVal);
	return retVal;
}

bool SimpleSocket::ioclt(long cmd, u_long* argp)
{
	static const char* functionName = "SimpleSocket::ioclt";
	
	int ret = ::ioctlsocket(socket_, cmd, argp);
	SHOW_SOCKET_ERROR((SOCKET_ERROR != ret), functionName, "::ioctlsocket");
	return (SOCKET_ERROR != ret);
}

struct sockaddr_in SimpleSocket::InitializeSockaddr(unsigned short port, const char* ipAddress)
{
	struct sockaddr_in retVal;
	retVal.sin_family = AF_INET;
	retVal.sin_addr.s_addr = inet_addr(ipAddress);  // dotted number form
	retVal.sin_port = htons(port);
	return retVal;
}

}
